# src/tools/mcp_client.py
from langchain_mcp_adapters.client import ClientSession, SSEConnection
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_mcp_adapters.sessions import create_session
from langchain_core.tools import BaseTool
from typing import List, Dict, Any, Optional
import asyncio


class MCPClient:
    """
    Unified MCP Client for connecting to MCP servers and executing tools.
    
    Provides a consistent interface for tool discovery and execution.
    """
    
    def __init__(self, server_url: str, api_key: Optional[str] = None):
        """
        Initialize MCP client.
        
        Args:
            server_url: MCP server SSE endpoint URL
            api_key: Optional authentication key
        """
        self.server_url = server_url
        self.api_key = api_key
        self._session: Optional[ClientSession] = None
        self._connection: Optional[SSEConnection] = None
        self._tools: List[BaseTool] = []
        self._is_connected = False
        self._session_context = None
    
    @property
    def is_connected(self) -> bool:
        """Check if client is connected to server"""
        return self._is_connected and self._session is not None
    
    @property
    def tools(self) -> List[BaseTool]:
        """Get available tools as LangChain BaseTool objects"""
        return self._tools
    
    async def connect(self) -> None:
        """
        Connect to MCP server and load available tools.
        
        Raises:
            ConnectionError: If connection fails
        """
        try:
            # Create connection configuration according to SSEConnection structure
            connection_config: SSEConnection = {
                "transport": "sse",
                "url": self.server_url
            }
            
            # Add headers if api_key is provided
            if self.api_key:
                connection_config["headers"] = {
                    "Authorization": f"Bearer {self.api_key}"
                }
            
            self._connection = connection_config
            
            # Create session using the context manager properly
            self._session_context = create_session(self._connection)
            self._session = await self._session_context.__aenter__()
            await self._session.initialize()
            
            # Load tools
            self._tools = await load_mcp_tools(self._session)
            self._is_connected = True
            
        except Exception as e:
            self._is_connected = False
            # Clean up if connection failed
            if self._session_context:
                try:
                    await self._session_context.__aexit__(None, None, None)
                except:
                    pass
                self._session_context = None
                self._session = None
            raise ConnectionError(f"Failed to connect to MCP server: {e}")
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """
        Execute a tool on the MCP server.
        
        Args:
            tool_name: Name of tool to execute
            arguments: Tool arguments
            
        Returns:
            Tool execution result
            
        Raises:
            RuntimeError: If client not connected
            ValueError: If tool not found
        """
        if not self.is_connected:
            raise RuntimeError("Client not connected. Call connect() first.")
        
        # Find and validate tool
        tool_names = [tool.name for tool in self._tools]
        if tool_name not in tool_names:
            raise ValueError(f"Tool '{tool_name}' not found. Available: {tool_names}")
        
        return await self._session.call_tool(tool_name, arguments)
    
    async def list_tools(self) -> List[Dict[str, str]]:
        """
        Get list of available tools with descriptions.
        
        Returns:
            List of tool information dictionaries
            
        Raises:
            RuntimeError: If client not connected
        """
        if not self.is_connected:
            raise RuntimeError("Client not connected. Call connect() first.")
        
        return [
            {
                "name": tool.name,
                "description": tool.description,
            }
            for tool in self._tools
        ]
    
    async def close(self) -> None:
        """Close connection and cleanup resources."""
        if self._session_context:
            try:
                await self._session_context.__aexit__(None, None, None)
            except:
                pass
        
        self._is_connected = False
        self._session = None
        self._session_context = None
        self._tools = []
    
    async def __aenter__(self):
        """Context manager entry - auto connect"""
        await self.connect()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - auto close"""
        await self.close()


# Factory function
async def create_mcp_client(server_url: str, api_key: Optional[str] = None) -> MCPClient:
    """
    Create and connect MCP client.
    
    Args:
        server_url: MCP server SSE endpoint URL
        api_key: Optional authentication key
        
    Returns:
        Connected MCPClient instance
    """
    client = MCPClient(server_url, api_key)
    await client.connect()
    return client


# Connection testing
async def test_connection(server_url: str, api_key: Optional[str] = None) -> Dict[str, Any]:
    """
    Test MCP server connection.
    
    Args:
        server_url: Server URL to test
        api_key: Optional API key for authentication
        
    Returns:
        Test results with status and tool information
    """
    client = None
    try:
        client = MCPClient(server_url, api_key)
        await client.connect()
        tools = await client.list_tools()
        
        result = {
            "success": True,
            "server_url": server_url,
            "tools_count": len(tools),
            "tools": [tool["name"] for tool in tools],
        }
        
    except Exception as e:
        result = {
            "success": False,
            "server_url": server_url,
            "error": str(e),
        }
    finally:
        if client:
            await client.close()
    
    return result


# Backward compatibility functions
async def create_mcp_connection(server_url: str, api_key: Optional[str] = None) -> MCPClient:
    """
    Create an MCP connection (backward compatibility).
    
    Args:
        server_url: The URL of the MCP server
        api_key: Optional API key for authentication
        
    Returns:
        MCPClient instance
    """
    client = MCPClient(server_url, api_key)
    await client.connect()
    return client


async def send_mcp_message(client: MCPClient, message_content: Dict[str, Any]) -> Dict[str, Any]:
    """
    Wrapper function to send messages (backward compatibility).
    
    Args:
        client: The MCP client instance
        message_content: The content of the message to send
        
    Returns:
        Response from the MCP server
    """
    # For backward compatibility, we'll treat this as a tool call if it contains tool info
    if "tool_name" in message_content and "arguments" in message_content:
        try:
            result = await client.call_tool(
                message_content["tool_name"], 
                message_content["arguments"]
            )
            return {"status": "success", "result": result}
        except Exception as e:
            return {"status": "error", "error": str(e)}
    else:
        # If it's just a message, we don't have a direct equivalent in this implementation
        return {"status": "error", "error": "Direct messaging not supported in this implementation"}


# Usage examples
if __name__ == "__main__":
    async def demo():
        """Demonstrate MCP client usage"""
        
        # Test connection
        test_url = "http://localhost:9009/sse"
        result = await test_connection(test_url)
        
        if result["success"]:
            print(f"✅ Connected to {result['server_url']}")
            print(f"🛠️  Found {result['tools_count']} tools: {', '.join(result['tools'])}")
        else:
            print(f"❌ Connection failed: {result['error']}")
            return
        
        # Example usage with manual connection
        print("\n--- Using manual connection ---")
        client = MCPClient(test_url)
        try:
            await client.connect()
            print(f"Connected with {len(client.tools)} tools")
            
            # List tools
            if client.tools:
                tool_info = await client.list_tools()
                print("Tool details:")
                for tool in tool_info:
                    print(f"  - {tool['name']}: {tool['description']}")
                    
        except Exception as e:
            print(f"Error during operation: {e}")
        finally:
            await client.close()
            print("Connection closed")
    
    asyncio.run(demo())